
import config
import sys
import pandas as pd
from fbprophet import Prophet
from fbprophet.plot import plot_plotly, plot_components_plotly
import json


def data_preparse(file):
    data = pd.read_csv(file)
    data = data.sort_values(by='trade_date', ascending=True)
    # print(data.tail())
    data = data[data['trade_date'] < 20210331]
    # print(data.tail())
    data = data.rename(columns={"trade_date": "ds"})
    print(data.head(10))
    x = [str(x)[0:4] + '-' + str(x)[4:6] + "-" + str(x)[6:8] for x in data['ds']]
    # print(x[:5])
    data['ds'] = x
    jine = pd.DataFrame(data, columns=['open', 'close', 'high', 'low'])
    jine['y'] = jine.apply(lambda x: sum(x) / 4, axis=1)
    # print(jine[:5])
    data['y'] = jine['y']
    return data


def model_file(file):
    data = data_preparse(file)
    print(data[:5])

    m = Prophet(daily_seasonality=True)
    m.fit(data)
    future = m.make_future_dataframe(periods=90)
    # print(future.tail(5))
    ret = m.predict(future)
    print(ret.tail(10))
    fig = plot_plotly(m, ret)
    fig.show()


if __name__ == '__main__':
    file = sys.argv[1]
    model_file(file)
